// ================================================
// CodeGeneratorDagster.ts
// ================================================

import {
  PipelineService, Node, Flow
} from './PipelineService';
import { BaseCodeGenerator, NodeObject } from './BaseCodeGenerator';

export class CodeGeneratorDagster extends BaseCodeGenerator {
  static generateDagsterCode(
    pipelineJson: string, 
    commands: any, 
    componentService: any, 
    variablesAutoNaming: boolean
  ): string {
    console.log("Inside generateDagsterCode method");

    const flow = PipelineService.filterPipeline(pipelineJson);
    const { nodesToTraverse, nodesMap, nodeDependencies } = this.computeNodesToTraverse(
      flow,
      'none',
      componentService
    );

    const dagsterImports = [
      'import dagster',
      'from dagster import op, job, Out, In, Nothing'
    ];

    const envVariablesCode = this.getEnvironmentVariableCode(pipelineJson, componentService);
    const connectionsCode = this.getConnectionCode(pipelineJson, componentService);

    const opDefinitions: string[] = [];
    const uniqueImports = new Set<string>();
    const uniqueDependencies = new Set<string>();
    const uniqueFunctions = new Set<string>();

    // Collect
    for (const nodeId of nodesToTraverse) {
      const node = nodesMap.get(nodeId);
      if (!node) continue;
      const config: any = node.data;
      const component = componentService.getComponent(node.type);
      component.provideImports({ config }).forEach(imp => uniqueImports.add(imp));
      if (typeof component.provideDependencies === 'function') {
        component.provideDependencies({ config }).forEach(d => uniqueDependencies.add(d));
      }
      if (typeof component.provideFunctions === 'function') {
        component.provideFunctions({ config }).forEach(f => uniqueFunctions.add(f));
      }
    }

    // Create ops (naive example)
    for (const nodeId of nodesToTraverse) {
      const node = nodesMap.get(nodeId);
      if (!node) continue;

      const config: any = node.data;
      const component = componentService.getComponent(node.type);
      const componentType = component._type;
      let opName = this.generateReadableName(config.customTitle || node.type);

      let opInputs: string[] = [];
      let opOutputs: string[] = [];
      let opCode = '';

      // Very naive logic (expand as needed)
      if (/processor/.test(componentType) || componentType.includes('to_documents')) {
        opInputs.push('input_data');
        opOutputs.push('result');
        const originalCode = component.generateComponentCode({
          config,
          inputName: 'input_data',
          outputName: 'result'
        });
        opCode = originalCode;
      } else if (componentType.includes('input')) {
        opOutputs.push('result');
        const originalCode = component.generateComponentCode({
          config,
          outputName: 'result'
        });
        opCode = originalCode;
      } else if (componentType.includes('output')) {
        opInputs.push('input_data');
        const originalCode = component.generateComponentCode({
          config,
          inputName: 'input_data'
        });
        opCode = originalCode;
      }

      // Build the definition
      const opDef = `
@op
def ${opName}(${opInputs.join(', ')}):
    """ ${config.customTitle || node.type} """
    ${opCode.split('\n').join('\n    ')}
    ${opOutputs.length > 0 ? 'return result' : 'return'}
`;
      opDefinitions.push(opDef);
    }

    // Build job definition
    const jobDefinition = this.generateDagsterJobDefinition(
      flow, 
      nodesMap, 
      nodeDependencies, 
      nodesToTraverse
    );

    // Combine
    const now = new Date();
    const dateString = now.toISOString().replace(/T/, ' ').replace(/\..+/, '');
    const dateComment = `# Source code generated by Amphi for Dagster\n# Date: ${dateString}\n`;
    const additionalDeps = `# Additional dependencies: dagster, ${Array.from(uniqueDependencies).join(', ')}\n`;

    const dagsterCode = [
      dateComment,
      additionalDeps,
      dagsterImports.join('\n'),
      '\n',
      Array.from(uniqueImports).join('\n'),
      '\n',
      envVariablesCode,
      '\n',
      connectionsCode,
      ...Array.from(uniqueFunctions),
      ...opDefinitions,
      jobDefinition
    ].join('');

    return this.formatVariables(dagsterCode);
  }

  static generateDagsterJobDefinition(
    flow: Flow,
    nodesMap: Map<string, Node>,
    nodeDependencies: Map<string, string[]>,
    nodesToTraverse: string[]
  ): string {
    const nodeToOp = new Map<string, string>();
    const processedNodes = new Set<string>();
    const resultVar = new Map<string, string>();

    for (const nodeId of nodesToTraverse) {
      const node = nodesMap.get(nodeId);
      if (!node) continue;
      const config: any = node.data;
      const opName = this.generateReadableName(config.customTitle || node.type);
      nodeToOp.set(nodeId, opName);

      const rVar = opName.replace('Op','Result');
      resultVar.set(nodeId, rVar);
    }

    let jobCode = '\n\n@job\ndef dagster_pipeline():\n';
    const dependencyGraph = new Map<string, string[]>();
    for (const edge of flow.edges) {
      if (!nodeToOp.has(edge.source) || !nodeToOp.has(edge.target)) continue;
      if (!dependencyGraph.has(edge.target)) {
        dependencyGraph.set(edge.target, []);
      }
      dependencyGraph.get(edge.target)!.push(edge.source);
    }

    const startingNodes = nodesToTraverse.filter(id => 
      !dependencyGraph.has(id) || dependencyGraph.get(id)!.length === 0
    );
    for (const nodeId of startingNodes) {
      const op = nodeToOp.get(nodeId);
      const r = resultVar.get(nodeId);
      if (!op || !r) continue;
      jobCode += `    ${r} = ${op}()\n`;
      processedNodes.add(nodeId);
    }

    for (const nodeId of nodesToTraverse) {
      if (processedNodes.has(nodeId)) continue;
      const op = nodeToOp.get(nodeId);
      const r = resultVar.get(nodeId)!;
      const deps = dependencyGraph.get(nodeId) || [];

      if (!op) continue;
      if (deps.length === 0) {
        jobCode += `    ${r} = ${op}()\n`;
      } else if (deps.length === 1) {
        const src = resultVar.get(deps[0])!;
        jobCode += `    ${r} = ${op}(${src})\n`;
      } else {
        const srcList = deps.map(d => resultVar.get(d)!).join(', ');
        jobCode += `    ${r} = ${op}(${srcList})\n`;
      }
      processedNodes.add(nodeId);
    }
    return jobCode;
  }

  static generateReadableName(rawName: string): string {
    const camelCaseName = rawName
      .split(/(?=[A-Z])/)
      .map((word, index) =>
        index === 0
          ? word.toLowerCase()
          : word.charAt(0).toUpperCase() + word.slice(1).toLowerCase()
      )
      .join('');
    return camelCaseName + 'Op';
  }
}